import redis.clients.jedis.Jedis;

public class TokenBucket {

    private Jedis jedis;
    private String bucketKey;
    private long rate; // 生成令牌的速率，以令牌/秒为单位
    private long lastRefillTime;

    public TokenBucket(String host, int port, String bucketKey, long capacity, long rate) {
        jedis = new Jedis(host, port);
        this.bucketKey = bucketKey;
        this.rate = rate;
        this.lastRefillTime = System.currentTimeMillis();

        // 初始化令牌桶，向Redis List中添加令牌
        jedis.del(bucketKey);
        for (int i = 0; i < capacity; i++) {
            jedis.rpush(bucketKey, "token");
        }
    }

    public boolean allowRequest() {
        long currentTime = System.currentTimeMillis();
        // 计算需要添加的令牌数量
        long elapsed = currentTime - lastRefillTime;
        long tokensToAdd = elapsed * rate / 1000; // 每秒生成的令牌数量  1000*2/1000 = 2
        System.out.println(tokensToAdd);
        // 获取当前令牌数量
        long currentTokens = jedis.llen(bucketKey);
        // 调整令牌数量，每秒追加2个
        if (tokensToAdd > 0) {
            for (int i = 0; i < tokensToAdd; i++) {
                jedis.rpush(bucketKey, "token");
            }
            lastRefillTime = currentTime;
        }
        // 如果令牌数量足够，允许请求 当前redis中存储的总数
        if (currentTokens > 0) {
            jedis.lpop(bucketKey); // 移除一个令牌 移出的是第一个令牌
            return true;
        } else {
            // 令牌不足，拒绝请求 写逻辑
            return false;
        }
    }

    public static void main(String[] args) {

        TokenBucket tokenBucket = new TokenBucket("192.168.7.81", 6379, "my_bucket", 10, 1); // 每秒生成2个令牌 ，令牌生成速率

        // 模拟请求
        for (int i = 0; i < 15; i++) {
            if (tokenBucket.allowRequest()) {
                System.out.println("Request " + i + " allowed.");
            } else {
                System.out.println("Request " + i + " denied.");
            }

            try {
                Thread.sleep(1000); // 模拟请求之间的延迟

            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

}
